"""
主题: 让对象支持上下文管理协议
问题: 你想让你的对象支持上下文管理协议(with语句)。
提示 : 
"""

from socket import socket, AF_INET, SOCK_STREAM


class LazyConnection:
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.sock = None

    def __enter__(self):
        if self.sock is not None:
            raise RuntimeError('Already connected')
        self.sock = socket(self.family, self.type)
        self.sock.connect(self.address)
        return self.sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.sock.close()
        self.sock = None

class LazyConnection2:
    '''嵌套使用连接'''
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.connections = []

    def __enter__(self):
        sock = socket(self.family, self.type)
        sock.connect(self.address)
        self.connections.append(sock)
        return sock

    def __exit__(self, exc_ty, exc_val, tb):
        self.connections.pop().close()

def recipe1():
    from functools import partial

    conn = LazyConnection(('www.python.org', 80))
    # Connection closed
    with conn as s:
        # conn.__enter__() executes: connection open
        s.send(b'GET /index.html HTTP/1.0\r\n')
        s.send(b'Host: www.python.org\r\n')
        s.send(b'\r\n')
        resp = b''.join(iter(partial(s.recv, 8192), b''))
        print(resp)
        # conn.__exit__() executes: connection closed

def recipe2():
    # Example use
    from functools import partial

    conn = LazyConnection2(('www.python.org', 80))
    with conn as s1:
        print(f"{s1 = }")
        with conn as s2:
            print(f"{s2 = }")
        # s1 and s2 are independent sockets

def main():
    print('recipe1'.center(20, '*'))
    recipe1()
    print('recipe2'.center(20, '*'))
    recipe2()



if __name__ == '__main__':
    main()           